# -*- coding: utf-8 -*-
# @Time    : 2023/9/29 10:36
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: UNite

import paddle
from paddle import nn
import paddle.nn.functional as F


def autopad(k, p=None):  # kernel, padding
    # Pad to 'same'
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p


def eval_act(act):
    if isinstance(act, str):
        return eval(act)
    elif act is True:
        return nn.Silu()
    elif isinstance(act, nn.Layer):
        return act
    else:
        return nn.Identity()


class AxialDWConv(nn.Layer):
    """ x` = x + DW_1xn(x) + DW_nx1(x) """

    def __init__(self, in_channels, ksz=7, dsz=1) -> None:
        super().__init__()
        psz = autopad(ksz) * dsz
        self.dw1 = nn.Conv2D(in_channels, in_channels, kernel_size=(ksz, 1), padding=(psz, 0), dilation=dsz, groups=in_channels)
        self.dw2 = nn.Conv2D(in_channels, in_channels, kernel_size=(1, ksz), padding=(0, psz), dilation=dsz, groups=in_channels)

    def forward(self, x):
        return x + self.dw1(x) + self.dw2(x)


class MyAxialDWConv(nn.Layer):
    """ x` = x + DW_1xn(DW_nx1(x)) """

    def __init__(self, in_channels, ksz=7, dsz=1) -> None:
        super().__init__()
        psz = autopad(ksz) * dsz
        self.hw = nn.Sequential(
            nn.Conv2D(in_channels, in_channels, kernel_size=(ksz, 1), padding=(psz, 0), dilation=dsz, groups=in_channels),
            nn.Conv2D(in_channels, in_channels, kernel_size=(1, ksz), padding=(0, psz), dilation=dsz, groups=in_channels)
        )
        self.wh = nn.Sequential(
            nn.Conv2D(in_channels, in_channels, kernel_size=(1, ksz), padding=(0, psz), dilation=dsz, groups=in_channels),
            nn.Conv2D(in_channels, in_channels, kernel_size=(ksz, 1), padding=(psz, 0), dilation=dsz, groups=in_channels)
        )

    def forward(self, x):
        return x + self.hw(x) + self.wh(x)


class AxialDPConv(nn.Layer):
    """ Axial depthwise and pointwise convolution """

    def __init__(self, in_channels, out_channels, ksz=7, act=True, use_my=False) -> None:
        super().__init__()
        self.dpconv = nn.Sequential(
            MyAxialDWConv(in_channels, ksz, 1) if use_my else AxialDWConv(in_channels, ksz, 1),
            nn.Conv2D(in_channels, out_channels, kernel_size=(1, 1)),
            eval_act(act)
        )

    def forward(self, x):
        return self.dpconv(x)


class AxialDWBottleNeck(nn.Layer):
    """ BottleNeck implement by AxialDWConv """

    def __init__(self, in_channels, ksz=3, act=True, use_my=False) -> None:
        super().__init__()
        assert in_channels % 4 == 0
        mid_channels = in_channels // 4

        self.fore = nn.Conv2D(in_channels, mid_channels, kernel_size=(1, 1))
        self.adw1 = MyAxialDWConv(mid_channels, ksz, 1) if use_my else AxialDWConv(mid_channels, ksz, 1)
        self.adw2 = MyAxialDWConv(mid_channels, ksz, 2) if use_my else AxialDWConv(mid_channels, ksz, 2)
        self.adw3 = MyAxialDWConv(mid_channels, ksz, 3) if use_my else AxialDWConv(mid_channels, ksz, 3)
        self.post = nn.Sequential(
            # nn.BatchNorm2D(in_channels),
            nn.Conv2D(in_channels, in_channels, kernel_size=(1, 1)),
            eval_act(act)
        )

    def forward(self, x):
        x0 = self.fore(x)
        x1 = self.adw1(x0)
        x2 = self.adw2(x0)
        x3 = self.adw3(x0)
        x4 = paddle.concat([x0, x1, x2, x3], axis=1)
        out = self.post(x4)
        return out


class ULiteNet(nn.Layer):
    """ Implementation of "1M PARAMETERS ARE ENOUGH? A LIGHTWEIGHT CNN-BASED MODEL FOR MEDICAL IMAGE SEGMENTATION" """

    def __init__(self, configs) -> None:
        super(ULiteNet, self).__init__()
        in_channels = configs["in_channels"] if "in_channels" in configs.keys() else 3
        mid_channels = configs["mid_channels"] if "mid_channels" in configs.keys() else (16, 32, 64, 128, 256, 512)
        num_classes = configs["num_classes"] if "num_classes" in configs.keys() else 3
        act = configs["act"] if "act" in configs.keys() else True
        use_my = configs["use_my"] if "use_my" in configs.keys() else False

        self.inconv = AxialDPConv(in_channels, mid_channels[0], 7, act, use_my)

        self.nc = len(mid_channels) - 1
        self.encoders = nn.Sequential(*(self.make_encoder(mid_channels[i], mid_channels[i + 1], act, use_my) for i in range(self.nc)))
        self.bottle = AxialDWBottleNeck(mid_channels[self.nc], 3, act, use_my)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.decoders = nn.Sequential(
            *(self.make_decoder(mid_channels[i + 1] + mid_channels[i], mid_channels[i], act, use_my) for i in range(self.nc)))

        # out_channels = 1 if num_classes <= 2 else num_classes
        self.outconv = nn.Conv2D(mid_channels[0], num_classes, kernel_size=(1, 1))

    def make_encoder(self, in_channels, out_channels, act, use_my):
        return nn.Sequential(
            nn.Sequential(
                MyAxialDWConv(in_channels, 7, 1) if use_my else AxialDWConv(in_channels, 7, 1),
                # nn.BatchNorm2D(in_channels)
            ),
            nn.Sequential(
                nn.Conv2D(in_channels, out_channels, kernel_size=(1, 1)),
                nn.MaxPool2D(kernel_size=(2, 2)),
                eval_act(act)
            )
        )

    def make_decoder(self, in_channels, out_channels, act, use_my):
        return nn.Sequential(
            nn.Conv2D(in_channels, out_channels, kernel_size=(1, 1)),
            # nn.BatchNorm2D(out_channels),
            MyAxialDWConv(out_channels, 7, 1) if use_my else AxialDWConv(out_channels, 7, 1),
            nn.Conv2D(out_channels, out_channels, kernel_size=(1, 1)),
            eval_act(act)
        )

    def forward(self, x):
        x0 = self.inconv(x)

        feats = []
        for i in range(self.nc):
            x0 = self.encoders[i][0](x0)
            feats.append(x0)
            x0 = self.encoders[i][1](x0)

        x0 = self.bottle(x0)
        for i in range(self.nc - 1, -1, -1):
            x0 = paddle.concat([self.up(x0), feats[i]], axis=1)
            x0 = self.decoders[i](x0)

        out = self.outconv(x0)
        return [out]


if __name__ == '__main__':
    model = ULiteNet({"use_my": True})
    paddle.flops(model, [1, 3, 512, 512])
